from app import db
import pickle
import re
from flask import current_app
import base64
import os
import random
from app.exceptions import MsgError
from app.parsers import random_str,compare_json_structure
import copy
import time


class Record(db.Model):
    id = db.Column(db.Integer,primary_key=True,autoincrement=True)
    belong_to_project = db.Column(db.Integer,db.ForeignKey('project.id'))
    risk_level = db.Column(db.Integer,default=0)
    checked = db.Column(db.Integer,default=0)
    timestamp = db.Column(db.Integer)
    req_info = db.Column(db.PickleType)

    @classmethod
    def create_record_by_json(cls,data):
        record = Record(belong_to_project=data['belong_to_project'],risk_level=data['risk_level'],
        checked=data['checked'],timestamp=data['timestamp'],req_info=data['req_info'])
        db.session.add(record)
        db.session.commit()
        return record

    @classmethod
    def set_state(cls,data):
        if 'record_id' not in data:
            raise MsgError("请求字段出错！")
        record = Record.query.get(data['record_id'])
        if record is None:
            raise MsgError("不存在的记录！")
        if 'checked' in data:
            record.checked = data['checked']
        if 'risk_level' in data:
            record.risk_level = data['risk_level']
        db.session.commit()

    @classmethod
    def delete_record_by_id(cls,record_id):
        record = Record.query.get(record_id)
        if record is None:
            raise MsgError("不存在的记录！")
        db.session.delete(record)
        db.session.commit()

    @classmethod
    def delete_record_by_id_list(cls,id_list):
        records = []
        if len(id_list) < 1:
            raise MsgError('指定了空列表！')
        for id in id_list:
            record = Record.query.get(id)
            if record is None:
                raise MsgError("不存在的记录！")
            records.append(record)
        for record in records:
            db.session.delete(record)
        db.session.commit()

    @classmethod
    def query_record(cls,req):
        ret = []
        record = Record.query
        if 'record_id' in req:
            record = record.get(req['record_id'])
            if record is not None:
                ret.append(record.to_json())
            return ret

        count = 10
        if 'count' in req and req['count'] > 1:
            count = req['count']
        if 'risk_level' in req:
            record = record.filter_by(risk_level=req['risk_level'])
        if 'checked' in req:
            record = record.filter_by(checked=req['checked'])
        if 'belong_to_project' in req:
            record = record.filter_by(belong_to_project=req['belong_to_project'])
        if 'timestamp_start' in req:
            record = record.filter(Record.timestamp>=req['timestamp_start']).order_by(Record.timestamp)
        if 'timestamp_end' in req:
            record = record.filter(Record.timestamp<=req['timestamp_end']).order_by(Record.timestamp)
        
        if 'timestamp_end' in req and 'timestamp_start' not in req:
            record = record.order_by(Record.timestamp.desc()).limit(count)
        else:
            record = record.limit(count)

        record = record.all()
        for r in record:
            ret.append(r.to_json())
        return ret

    def to_json(self):
        ret = {}
        ret['record_id'] = self.id
        ret['belong_to_project'] = self.project.id
        ret['risk_level'] = self.risk_level
        ret['checked'] = self.checked
        ret['timestamp'] = self.timestamp
        ret['req_info'] = self.req_info
        return ret

class Project(db.Model):
    id = db.Column(db.Integer,primary_key=True,autoincrement=True)
    title = db.Column(db.String(225))
    description = db.Column(db.String(1024))
    config = db.Column(db.PickleType)
    records = db.relationship('Record',backref="project",lazy='dynamic')
    code_index = db.Column(db.String(16),unique=True) # 指定存储payload的文件名

    def get_project_url(self): # 获取项目的真实地址
        return "//{}/xss/{}".format(current_app.config.get('WEB_ROOT'),self.config['address'])

    @classmethod
    def get_all_project_info_with_json(cls):
        ret = []
        projects = Project.query.all()
        for p in projects:
            tmp = p.to_json()
            tmp['url'] = p.get_project_url()
            tmp.pop('config')
            tmp['record_count'] = p.records.count()
            ret.append(tmp)
        return ret

    @classmethod
    def get_project_detail_by_id(cls,project_id):
        project = Project.query.get(project_id)
        if project is None:
            raise MsgError('项目不存在！')
        ret = project.to_json()
        return ret

    @classmethod
    def delete_project(cls,project_id):
        project = Project.query.get(project_id)
        if project is None:
            raise MsgError('项目不存在！')
        if(project.records.count() > 0):
            raise MsgError('项目“{}”存在访问记录，无法删除！'.format(project.title))
        project.delete_code()
        db.session.delete(project)
        db.session.commit()

    @classmethod
    def delete_project_by_id_list(cls,id_list):
        if len(id_list)<1:
            raise MsgError('指定了空列表！')
        projects = []
        for project_id in id_list:
            project = Project.query.get(project_id)
            if project is None:
                raise MsgError('项目不存在！')
            projects.append(project)
        for project in projects:
            if project.records.count() > 0:
                raise MsgError('项目“{}”存在访问记录，无法删除！'.format(project.title))
        for project in projects:
            project.delete_code()
            db.session.delete(project)
        db.session.commit()

    def create_code(self):# 创建项目对应的payload代码
        tmp_file = os.path.join(current_app.config.get('PROJECT_CODE_DIR'),self.code_index)
        with open(tmp_file,'wb') as f:
            data = base64.b64decode(self.config['extra_code'].encode('utf8')).decode('utf8',errors='ignore')+'\r\n\r\n'
            for i in self.config['payload_list']:
                payload = Payload.query.get(i['payload_id'])
                if payload is None:
                    raise MsgError('指定的payload不存在！')
                data += payload.get_payload_with_options(i['options']) + '\r\n\r\n'
            f.write(data.encode('utf8'))

    def delete_code(self):
        tmp_file = os.path.join(current_app.config.get('PROJECT_CODE_DIR'),self.code_index)
        try:
            os.remove(tmp_file)
        except:
            raise MsgError('系统错误！无法删除项目对应文件！')


    @classmethod
    def get_code_by_project_id(cls,project_id):
        project = Project.query.get(project_id)
        if project is None:
            raise MsgError('不存在的项目！')
        return project.get_code()

    def get_code(self):# 获取项目对应的payload代码
        tmp_file = os.path.join(current_app.config.get('PROJECT_CODE_DIR'),self.code_index)
        data = ''
        with open(tmp_file,'rb') as f:
            data = f.read()
        return data.decode('utf8',errors='ignore')


    @classmethod
    def get_project_by_address(cls,address): 
        projects = Project.query.all()
        ret = None
        for p in projects:
            if p.config['address'] == address:
                ret = p
                break
        return ret

    @classmethod
    def addressIsOK(cls,address):
        good_key = '0-9a-zA-Z'
        if len(re.findall(r'^[{}]+$'.format(good_key),address)) < 1:
            raise MsgError("address存在不可用的字符，可用字符有{}".format(good_key))

        projects = Project.query.all()
        for project in projects:
            if project.config['address'] == address:
                raise MsgError("提交的address已存在，请使用其他值！")

    @classmethod
    def create_or_modify_by_json(cls,data):
        if 'project_id' in data: # 修改
            project = Project.query.get(data['project_id'])
            if project is None:
                raise MsgError('不存在的项目！')
            if 'title' in data:
                project.title = data['title']
            if 'description' in data:
                project.description = data['description']
            if 'config' in data:
                config = data['config']
                new_config = copy.deepcopy(project.config)
                if 'address' in config:
                    Project.addressIsOK(config['address'])
                    new_config['address'] = config['address']
                if 'extra_response_header' in config:
                    new_config['extra_response_header'] = config['extra_response_header']
                if 'payload_list' in config:
                    new_config['payload_list'] = config['payload_list']
                if 'extra_code' in config:
                    new_config['extra_code'] = config['extra_code']
            project.config = new_config
            project.create_code()
            db.session.commit()
            return project
        else: # 创建
            for i in ['title','description','config']:
                if i not in data:
                    raise MsgError('参数缺失！')
            config = { # 用于结构比较，值不具有意义
                    "address": '/test_xss',
                    "extra_response_header": [ 
                        ['TestHeader','test xss!'], 
                        ['key','value']
                    ],
                    "payload_list": [ 
                        {
                            "payload_id": 1,
                            "options": [
                                ['return_ip','127.0.0.1']
                            ]
                        },
                        {
                            "payload_id": 3,
                            "options": [
                                ['return_ip','127.0.0.1']
                            ]
                        }
                    ], 
                    "extra_code": 'adfdafdsf='
            }
            if not compare_json_structure(config,data['config']):
                raise MsgError('请求参数错误！')
            Project.addressIsOK(data['config']['address'])
            project = Project(title=data['title'],description=data['description'],config=data['config'],
            code_index=random_str())
            project.create_code()
            db.session.add(project)
            db.session.commit()
            return project


    def to_json(self):
        ret = {}
        ret['title'] = self.title
        ret['description'] = self.description
        ret['project_id'] = self.id
        ret['config'] = self.config
        return ret

class Payload(db.Model):
    id = db.Column(db.Integer,primary_key=True,autoincrement=True)
    title = db.Column(db.String(225))
    description = db.Column(db.String(1024))
    category_id = db.Column(db.Integer,db.ForeignKey('category.id'))
    value_index = db.Column(db.String(16),unique=True) # 指定存储payload的文件名
    options = db.Column(db.PickleType) # 存储选项相关的json数据，如["return_ip",...],每一项参数名

    @classmethod
    def get_options(cls,text): # 从text文本中读取选项参数
        ret = []
        search_re = r"<x [a-zA-Z0-9_\$]+ x>"
        result = re.findall(search_re,text)
        result = [i[3:-3] for i in result]
        for i in set(result):
            ret.append(i)
        return ret

    def get_payload_with_options(self,options):# 返回设置了option后的payload代码
        data = self.get_payload_value().decode('utf8',errors='ignore')
        for i in options:
            data = data.replace('<x {} x>'.format(i[0]),i[1])
        return data

    def hasProjectNeed(self): # 判断是否有项目依赖于改payload
        projects =  Project.query.all()
        for project in projects:
            for p in project.config['payload_list']:
                if p['payload_id'] == self.id:
                    return True

    @classmethod
    def delete_payload(cls,payload_id):
        paylaod = Payload.query.get(payload_id)
        if paylaod is None:
            raise MsgError('payload不存在！')
        if paylaod.hasProjectNeed():
            raise MsgError('存在项目依赖于改payload，无法删除！')
        tmp_file = os.path.join(current_app.config.get('PAYLOAD_DIR'),paylaod.value_index)
        try:
            os.remove(tmp_file)
        except:
            raise MsgError('系统错误！无法删除payload文件！')
        db.session.delete(paylaod)
        db.session.commit()

    @classmethod
    def delete_payload_by_id_list(cls,id_list):
        if len(id_list)<1:
            raise MsgError('指定了空列表！')
        payloads = []
        for payload_id in id_list:
            paylaod = Payload.query.get(payload_id)
            if paylaod is None:
                raise MsgError('payload不存在！')
            payloads.append(paylaod)
        for payload in payloads:
            if paylaod.hasProjectNeed():
                raise MsgError('存在项目依赖于payload“{}”，无法删除！'.format(payload.title))
        for paylaod in payloads:
            tmp_file = os.path.join(current_app.config.get('PAYLOAD_DIR'),paylaod.value_index)
            try:
                os.remove(tmp_file)
            except:
                raise MsgError('系统错误！无法删除“{}”的payload文件！'.format(payload.title))
            db.session.delete(payload)
        db.session.commit()

    @classmethod
    def get_all_payload_with_json(cls):
        ret = []
        payload_list = Payload.query.all()
        for payload in payload_list:
            tmp = payload.to_json()
            tmp.pop('value')
            ret.append(tmp)
        return ret

    def get_payload_value(self): # 返回的是二进制数据，bytes类型
        tmp_file = os.path.join(current_app.config.get('PAYLOAD_DIR'),self.value_index)
        ret = b''
        with open(tmp_file,'rb') as f:
            ret = f.read()
        return ret

    def set_payload_value(self,data): # 设置payload，data为bytes类型
        tmp_file = os.path.join(current_app.config.get('PAYLOAD_DIR'),self.value_index)
        with open(tmp_file,'wb') as f:
            f.write(data)

    @classmethod
    def encode(cls,text): # text为bytes类型，对text进行编码或加密
        return base64.b64encode(text).decode('utf8',errors='ingore')

    @classmethod
    def decode(cls,text):
        return base64.b64decode(text.encode('utf8'))

    @classmethod
    def get_payload_info_with_json(cls,payload_id):
        payload = Payload.query.get(payload_id)
        if payload is None:
            raise MsgError('不存在的payload！')
        ret = payload.to_json()
        return ret

    def to_json(self):
        ret = {}
        ret['title'] = self.title
        ret['description'] = self.description
        ret['category_id'] = self.category.id
        ret['category'] = self.category.name
        ret['payload_id'] = self.id
        ret['value'] = self.encode(self.get_payload_value())
        ret['options'] = self.options

        return ret


    @classmethod
    def create_or_modify_by_json(cls,data): # 查询是否有正确的payload_id进行修改，否则创建一个新的payload
        columns = data.keys()
        if 'payload_id' in data: # 修改payload
            payload = cls.query.get(data['payload_id'])
            if payload is None:
                raise MsgError('不存在的payload！')
            if 'title' in columns:
                payload.title = data['title']
            if 'description' in columns:
                payload.description = data['description']
            if 'category_id' in columns:
                category = Category.query.get(data['category_id'])
                if category is None:
                    raise MsgError('指定的分类不存在！')
                payload.category = category
            if 'value' in columns:
                code = Payload.decode(data['value'])
                payload.set_payload_value(code)
                options = Payload.get_options(code.decode('utf8',errors='ignore'))
                payload.options = options
            db.session.commit()
            return payload
        else: # 新建一个payload
            for i in ['title','description','category_id','value']:
                if i not in columns:
                    raise MsgError('请求字段出错！')
            category = Category.query.get(data['category_id'])
            if category is None:
                raise MsgError('指定的分类不存在！')
            value_index = random_str()
            tmp_file = os.path.join(current_app.config.get('PAYLOAD_DIR'),value_index)
            with open(tmp_file,'wb') as f:
                f.write(Payload.decode(data['value']))
            options = Payload.get_options(Payload.decode(data['value']).decode('utf8',errors='ignore'))
            payload = Payload(title=data['title'],description=data['description'],category=category,value_index=value_index,\
                options=options)
            db.session.add(payload)
            db.session.commit()
            return payload

        
class Category(db.Model):
    id = db.Column(db.Integer,primary_key=True,autoincrement=True)
    name = db.Column(db.String(225),default='未命名分类')
    payloads = db.relationship('Payload',backref="category",lazy='dynamic')


    @classmethod
    def get_all_category_with_json(cls):
        ret = []
        category_list = Category.query.all()
        for category in category_list:
            tmp = category.to_json()
            ret.append(tmp)
        return ret

    @classmethod
    def rename(cls,category_id,name):
        category = Category.query.get(category_id)
        if category is None:
            raise MsgError("不存在的分类！")
        category.name = name
        db.session.commit()
        return category

    @classmethod
    def delete_category(cls,category_id):
        if category_id == 1:
            raise MsgError('该分类不可删除！')
        category = Category.query.get(category_id)
        if category is None:
            raise MsgError('分类不存在！')
        if category.payloads.count() > 0:
            raise MsgError("该分类下存在payload，无法删除！")
        db.session.delete(category)
        db.session.commit()

    def to_json(self):
        ret = {}
        ret['category_id'] = self.id
        ret['name'] = self.name
        ret['payload_list'] = [payload.id for payload in self.payloads]
        return ret

    @classmethod
    def create_category(cls,name):
        rel = Category.query.filter_by(name=name).first()
        if rel is not None:
            raise MsgError('已存在的分类名！')
        category = Category(name=name)
        db.session.add(category)
        db.session.commit()
        return category
        
class Session(db.Model):
    id = db.Column(db.Integer,primary_key=True,autoincrement=True)
    login_id = db.Column(db.String(225))
    create_timestamp = db.Column(db.Integer)

    @classmethod
    def login_check(cls,login_id):
        session = Session.query.filter_by(login_id=login_id).first()
        if session is None:
            return False
        if(session.session_check_timeout()):
            db.session.delete(session)
            db.session.commit()
            return False
        return True

    def session_check_timeout(self):
        timeout = current_app.config.get("SESSION_TIMEOUT")
        if int(time.time()) - self.create_timestamp > timeout:
            return True
        return False

    @classmethod
    def session_destroy(cls,login_id):
        session = Session.query.filter_by(login_id=login_id).first()
        if session is None:
            return False
        db.session.delete(session)
        db.session.commit()
        return True

    @classmethod
    def create_session(cls,password): 
        if password == current_app.config.get('LOGIN_PASSWORD'):
            session = Session(login_id=random_str(48),create_timestamp=int(time.time()))
            db.session.add(session)
            db.session.commit()
            return session
        else:
            raise MsgError("密码错误！")
